import requests
import functools
import asyncio
from concurrent.futures import ThreadPoolExecutor
import json 
import websocket
import random
from datetime import datetime, timedelta
from core import setting, mgClient
import pymongo as mg
from poe_api_wrapper import PoeApi
from aiohttp import web
import asyncio
import uuid
db = mgClient['chat']
draw_collection = db.get_collection('draw')
dall_collection = db.get_collection('dall')

def check_traffic(ip):
    """
    It checks the traffic of the given IP address.
    
    :param ip: The IP address of the host to check
    """
    now = datetime.now()
    if ip != "221.215.48.6":
        traffic = draw_collection.find_one({"ip": ip})  # 查找指定IP地址的流量信息
        if not traffic:
            traffic = {"ip": ip, "data": [{"timestamp": now, "count": 1}]}
            draw_collection.insert_one(traffic)  # 插入新的流量信息
            return False
        else:
            recent_count = sum([t["count"] for t in traffic["data"]])
            # 删除超过24小时的旧流量数据
            traffic["data"] = [t for t in traffic["data"] if now - t["timestamp"] <= timedelta(hours=24)]
            if recent_count >= 500:  # 设置流量限制为10分钟内最多60个请求
                return True
            else:
                traffic["data"].append({"timestamp": now, "count": 1})
                draw_collection.update_one({"ip": ip}, {"$set": traffic})  # 更新流量信息
                return False
    else:
        return False
def check_traffic_3(ip):
    """
    It checks the traffic of the given IP address.
    
    :param ip: The IP address of the host to check
    """
    now = datetime.now()
    if True:
        traffic = dall_collection.find_one({"ip": ip})  # 查找指定IP地址的流量信息
        if not traffic:
            traffic = {"ip": ip, "data": [{"timestamp": now, "count": 1}]}
            dall_collection.insert_one(traffic)  # 插入新的流量信息
            return False
        else:
            recent_count = sum([t["count"] for t in traffic["data"]])
            # 删除超过24小时的旧流量数据
            traffic["data"] = [t for t in traffic["data"] if now - t["timestamp"] <= timedelta(hours=24)]
            if recent_count >= 1:  # 设置流量限制为10分钟内最多60个请求
                return True
            else:
                traffic["data"].append({"timestamp": now, "count": 1})
                dall_collection.update_one({"ip": ip}, {"$set": traffic})  # 更新流量信息
                return False
    else:
        return False
def get_api_key(api_keys, api_key_index):
    """
    It returns the next API key in the list of API keys
    :return: The api_key is being returned.
    """
    # print(api_key_index)
    api_key = api_keys[api_key_index]
    api_key_index = (api_key_index + 1) % len(api_keys)
    return api_key, api_key_index

def get_api_keys(key_file_path):
    """
    It opens the file at the path specified by the argument key_file_path, reads each line of the
    file, strips the newline character from each line, and returns a list of the lines
    
    :param key_file_path: The path to the file containing the API keys
    :return: A list of API keys
    """
    api_keys = []
    with open(key_file_path, 'r') as f:
        for line in f:
            api_keys.append(line.strip())
    return api_keys
def get_url(message,api_key,model='dall-e-3',size='1024x1024',quality='standard'):
    
    key_values =[
        'your_key'
    ]
    api_key = random.choice(key_values)

    url = "https://api.openai.com/v1/images/generations"
    headers = {
        "Content-Type": "application/json",
        "Authorization": api_key
    }
    data = {
        "model": model,
        "prompt": message,
        "n": 1,
        "size": size,
        "quality": quality,
    }

    response = requests.post(url, headers=headers, json=data)
    if response.status_code != 200:
        print(response.status_code)
        print(response.text)
        raise Exception("Failed to generate image")
    response_data = response.json()

    generated_image_url = response_data["data"][0]["url"]
    return generated_image_url
async def generate_image_dall_e(message,api_key,model):
    mode_json = {
        '方':{
            'size':'1024x1024',
            'model':'dall-e-3',
            'quality':'standard'
        },
        '高':{
            'size':'1024x1792',
            'model':'dall-e-3',
            'quality':'standard'
        },
        '宽':{
            'size':'1792x1024',
            'model':'dall-e-3',
            'quality':'standard'
        },
        '16:9':{
            'size':'1792x1024',
            'model':'dall-e-3',
            'quality':'standard'
        },
        '9:16':{
            'size':'1024x1792',
            'model':'dall-e-3',
            'quality':'standard'
        },
        'dall-e-3-高质量':{
            'size':'1024x1024',
            'model':'dall-e-3',
            'quality':'hd'
        },
    }
    size = mode_json[model]['size']
    quality = mode_json[model]['quality']
    model = mode_json[model]['model']
    try:
        loop = asyncio.get_running_loop()
        response = await loop.run_in_executor(None,get_url, message,api_key,model,size,quality)
        # print(response)
        # url = response.split('private/')[-1].replace('/','goodpicture')
        return response
    except:
        return '安全检查没通过' 
def send_message(client, mode, message):
    for chunk in client.send_message(mode, message):
        print(chunk)
        # print(chunk["response"], end="", flush=True)
    response = chunk["response"]
    return response

async def generate_image_poe(client_token,message,mode):
    
    try:
        client = PoeApi(client_token)
        # print('adfadsff')
        loop = asyncio.get_running_loop()
        response = await loop.run_in_executor(None, send_message, client, mode, message)
        # print('adadasd')
        # ()内的是url
        url = response.split('\n')[0].split(']:')[-1]
        print(url)
        # url_task_id = url.split('/')[-1]
        return url 
    except Exception as e:
        # print(e)
        # client.delete_chat(mode, del_all=True)
        client.disconnect_ws()
        client.client.close()
        del client
        raise e
    finally:
        # client.delete_chat(mode, del_all=True)
        client.disconnect_ws()
        client.client.close()
        del client
async def async_request(prompt: str, api_key: str):
    
    data = {
        "model": "gpt-3.5-turbo",
        "messages": prompt,
        "temperature": 1,
        # "max_tokens": 4096,
        "stream": True,  # Enable streaming API
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer "+api_key,
    }
    url = "https://api.openai.com/v1/chat/completions"
    loop = asyncio.get_running_loop()
    request = functools.partial(requests.post, url, headers=headers, json=data, stream = True, timeout = 3)
    try:
        with ThreadPoolExecutor() as executor:
            response = await loop.run_in_executor(
                executor, request
            )
    except:
        response = requests.Response()
        response.status_code = 500
    return response


async def request_anytext(task_id, data,img_url):
    # 模拟一个耗时的任务
    url = "img_server_addresss/anytext"
    response = requests.post(url, json=data)
    with open(f'/root/backend_test/backend/DrawAPI/AnyText/{task_id}.jpeg', 'wb') as f:
        f.write(response.content)
    img_url[task_id] = f'/root/backend_test/backend/DrawAPI/AnyText/{task_id}.jpeg'


